# 新增路径编码脚本

import torch
import numpy as np
from svgpathtools import parse_path
import os
from src.data_preprocessing.svg_constants import SVG_COMMANDS

def sanitize_svg_path(svg_content):
    """保留所有SVG命令的路径清理"""
    if isinstance(svg_content, str) and os.path.exists(svg_content):
        with open(svg_content, 'r') as f:
            svg_content = f.read()
    
    # 提取路径数据
    if '<path' in svg_content:
        svg_content = svg_content.split('d="')[-1].split('"')[0]
    
    # 修改为使用SVG_COMMANDS常量
    for cmd in SVG_COMMANDS:
        svg_content = svg_content.replace(cmd, f' {cmd} ')
    
    return ' '.join(svg_content.split())  # 压缩多余空格

def svg_to_tensor(svg_path, img_size=128):
    """增强版SVG转张量"""
    try:
        clean_path = sanitize_svg_path(svg_path)
        path = parse_path(clean_path)
        
        # 修改初始化画布和绘制逻辑
        img = np.ones((img_size, img_size))  # 初始化为白底
        for segment in path:
            samples = 100
            for t in np.linspace(0, 1, samples):
                x, y = segment.point(t).real, segment.point(t).imag
                x = int(np.clip(x, 0.0, 1.0) * (img_size-1))  # 确保在图像范围内
                y = int(np.clip(y, 0.0, 1.0) * (img_size-1))
                img[y,x] = 0.0  # 绘制黑字
        
        return torch.from_numpy(img).float()  # 返回白底黑字
    except Exception as e:
        print(f"处理SVG文件失败: {svg_path}, 错误: {str(e)}")
        return None

def process_svg_batch(svg_metadata, output_tensor_dir):
    """批量处理SVG文件"""
    os.makedirs(output_tensor_dir, exist_ok=True)
    
    for record in svg_metadata:
        svg_path = os.path.join('data/processed/svg', record['svg_path'])
        tensor = svg_to_tensor(svg_path)
        torch.save(tensor, os.path.join(output_tensor_dir, f"{record['font']}_{record['char']}.pt"))

def process_font_directory(svg_dir, output_base_dir):
    """处理单个字体目录"""
    font_name = os.path.basename(svg_dir)
    output_dir = os.path.join(output_base_dir, font_name)
    os.makedirs(output_dir, exist_ok=True)
    
    svg_files = [f for f in os.listdir(svg_dir) if f.endswith(".svg")]
    metadata = [{
        "font": font_name,
        "char": f.split(".")[0],  # 假设文件名就是字符名
        "svg_path": os.path.join(svg_dir, f)
    } for f in svg_files]
    
    process_svg_batch(metadata, output_dir)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--svg_root", default="/Users/weiyin/Workspaces/TypusMotrix/data/processed/svg_paths")
    parser.add_argument("--output_root", default="/Users/weiyin/Workspaces/TypusMotrix/data/processed/tensors")
    args = parser.parse_args()

    # 处理所有字体子目录
    font_dirs = [d for d in os.listdir(args.svg_root) 
                if os.path.isdir(os.path.join(args.svg_root, d))]
    
    for font_dir in font_dirs:
        full_path = os.path.join(args.svg_root, font_dir)
        print(f"正在处理字体目录: {font_dir}")
        process_font_directory(full_path, args.output_root)
